import React, { useState, useEffect } from 'react';
import {
    Table,
    Card,
    Button,
    Modal,
    Input,
    message,
    Space,
    Popconfirm,
    Typography
} from 'antd';
import {
    PlusOutlined,
    DeleteOutlined,
    ReloadOutlined,
    ExclamationCircleOutlined,
} from '@ant-design/icons';

const { Title } = Typography;
const API_BASE_URL = import.meta.env.VITE_OLLAMA_URL;

const OllamaManager = () => {
    const [models, setModels] = useState([]);
    const [loading, setLoading] = useState(false);
    const [isModalOpen, setIsModalOpen] = useState(false);
    const [newModelName, setNewModelName] = useState('');
    const [messageApi, contextHolder] = message.useMessage();

    // 添加停止模型的方法
    const stopModel = async (modelName) => {
        try {
            const response = await fetch(`${API_BASE_URL}/api/generate`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify({
                    model: modelName,
                    prompt: "", // 空提示
                    keep_alive: 0 // 设置为0释放模型
                })
            });

            if (response.ok) {
                messageApi.success(`模型 ${modelName} 已停止`);
                fetchModels(); // 刷新列表
            } else {
                messageApi.error('停止模型失败');
            }
        } catch (err) {
            console.error('Error stopping model:', err);
            messageApi.error('停止模型失败');
        }
    };

    // 获取模型列表
    const fetchModels = async () => {
        setLoading(true);
        try {
            const response = await fetch(`${API_BASE_URL}/api/tags`);
            const data = await response.json();
            setModels(data.models || []);
        } catch (err) {
            console.error('Error fetching models:', err);
            messageApi.error('获取模型列表失败');
        } finally {
            setLoading(false);
        }
    };

    // 拉取新模型
    const pullModel = async () => {
        if (!newModelName) {
            messageApi.warning('请输入模型名称');
            return;
        }

        setLoading(true);
        try {
            messageApi.loading('正在下载模型请等待下载完成..');
            const response = await fetch(`${API_BASE_URL}/api/pull`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify({ name: newModelName }),
            });

            if (!response.ok) throw new Error('拉取模型失败');

            // 读取流数据
            const reader = response.body.getReader();
            let downloading = true;

            while (downloading) {
                const { value, done } = await reader.read();
                if (done) {
                    downloading = false;
                    break;
                }

                // 解析返回的进度数据
                const text = new TextDecoder().decode(value);
                const lines = text.split('\n').filter(line => line.trim());

                for (const line of lines) {
                    try {
                        const data = JSON.parse(line);
                        // 更新下载进度提示
                        if (data.status === 'downloading') {
                            messageApi.loading(`正在下载: ${Math.round(data.completed * 100)}%`);
                        } else if (data.status === 'completed') {
                            messageApi.success('模型下载完成！');
                            downloading = false;
                        }
                    } catch (e) {
                        console.error('解析进度数据失败:', e);
                    }
                }
            }

            // 所有数据接收完成后再关闭窗口
            setIsModalOpen(false);
            setNewModelName('');
            fetchModels();
        } catch (err) {
            console.error('Error pulling model:', err);
            messageApi.error('拉取模型失败');
        } finally {
            setLoading(false);
        }
    };

    // 删除模型
    const deleteModel = async (name) => {
        setLoading(true);
        try {
            const response = await fetch(`${API_BASE_URL}/api/delete`, {
                method: 'DELETE',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify({ name }),
            });

            if (!response.ok) throw new Error('删除模型失败');

            messageApi.success('删除成功');
            await fetchModels();
        } catch (err) {
            console.error('Error deleting model:', err);
            messageApi.error('删除模型失败');
        } finally {
            setLoading(false);
        }
    };

    useEffect(() => {
        fetchModels();
    }, []);

    const columns = [
        {
            title: '模型名称',
            dataIndex: 'name',
            key: 'name',
            render: (text) => <a>{text}</a>,
        },
        {
            title: '大小',
            dataIndex: 'size',
            key: 'size',
            render: (size) => `${(size / (1024 * 1024 * 1024)).toFixed(2)} GB`,
        },
        {
            title: '修改时间',
            dataIndex: 'modified_at',
            key: 'modified_at',
            render: (modified) => new Date(modified).toLocaleString(),
        },
        {
            title: '操作',
            key: 'action',
            width: 230,
            render: (_, record) => (
                <Space>
                    <Popconfirm
                        title="删除模型"
                        description={`确定要删除模型 ${record.name} 吗？`}
                        onConfirm={() => deleteModel(record.name)}
                        okText="确定"
                        cancelText="取消"
                        icon={<ExclamationCircleOutlined style={{ color: 'red' }} />}
                    >
                        <Button type="link" danger icon={<DeleteOutlined />}>
                            删除
                        </Button>
                    </Popconfirm>
                    <Button type="link" onClick={() => stopModel(record.name)}>停用</Button>
                </Space>
            ),
        },
    ];

    return (
        <div>
            {contextHolder}
            <Card
                title={<Title level={4}>模型管理</Title>}
                extra={
                    <Space>
                        <Button
                            icon={<ReloadOutlined />}
                            onClick={fetchModels}
                            loading={loading}
                        >
                            刷新
                        </Button>
                        <Button
                            type="primary"
                            icon={<PlusOutlined />}
                            onClick={() => setIsModalOpen(true)}
                        >
                            拉取模型
                        </Button>
                    </Space>
                }
            >
                <Table
                    columns={columns}
                    dataSource={models}
                    rowKey="name"
                    loading={loading}
                    pagination={false}
                    locale={{
                        emptyText: '暂无模型',
                    }}
                />
            </Card>

            <Modal
                title="拉取新模型"
                open={isModalOpen}
                onOk={pullModel}
                onCancel={() => {
                    setIsModalOpen(false);
                    setNewModelName('');
                }}
                confirmLoading={loading}
            >
                <Input
                    placeholder="输入模型名称 (例如: llama2)"
                    value={newModelName}
                    onChange={(e) => setNewModelName(e.target.value)}
                    style={{ marginTop: '16px' }}
                />
            </Modal>
        </div>
    );
};

export default OllamaManager;